"""
复合型序列(趋势、季节、周期和随机成分)：Holt-Winters季节性预测模型
"""
import numpy as np
import pandas as pd
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import TimeSeriesSplit
from statsmodels.tsa.api import ExponentialSmoothing
# from statsmodels.tsa.holtwinters import ExponentialSmoothing


def holt_winters_trend_season(series, window):
    """
    复合型序列(趋势、季节、周期和随机成分)：Holt-Winters季节性预测模型
    :param series:
    :param window:
    :return:
    """
    if isinstance(series, pd.Series):
        data_set = series.to_numpy()

    if not isinstance(series, np.ndarray):
        data_set = np.array(series)
    else:
        data_set = series

    rms_holt = []

    tscv = TimeSeriesSplit()
    for train_idx, test_idx in tscv.split(data_set):
        train, test = data_set[train_idx], data_set[test_idx]
        peroid = len(test_idx)

        # holt
        """        
        seasonal_periods：周期数：季度：4；月度：12；周：7
        trend：
          趋势成分类型：“add”, “mul”, “additive”, 
                         “multiplicative”, None
        seasonal：
          季节成分类型：“add”, “mul”, “additive”, 
                        “multiplicative”, None
        """
        holt_model = ExponentialSmoothing(data_set, seasonal_periods=window,
                                          trend='add', seasonal='add',).fit()
        y_pred = holt_model.forecast(peroid)
        rms_holt.append(np.sqrt(mean_squared_error(test, y_pred)))

    print('rms: {}'.format(np.mean(rms_holt)))
